import torch
import torch.nn as nn
import torchvision.models as models
from torchsummary import summary
class self_net(nn.Module):
    def __init__(self,
                 in_channels: int = 1,
                 num_classes: int = 4):
        super(self_net, self).__init__()
        vgg11 = models.vgg11_bn(pretrained=False)

        # 提取 VGG 的前几层（除去全连接层）
        self.features = vgg11.features

        # 构建上采样层（反卷积层或者 bilinear upsample）
        self.upsample = nn.ConvTranspose2d(512, num_classes, kernel_size=72, stride=32, padding=16, bias=False)

    def forward(self, x):
        # 提取 VGG 特征
        x = self.features(x)

        # 上采样到原始输入大小
        x = self.upsample(x)

        return x
if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = self_net().to(device)
    summary(model, (3, 200, 200))

